import torch.nn as nn

class Flatten(nn.Module):
  def __init__(self):
    super(Flatten, self).__init__()

  def forward(self, input):
    return input.view(input.size(0), -1)

class TestNet(nn.Module):
  def __init__(self):
    super(TestNet, self).__init__()
    self.net = nn.Sequential(nn.Conv2d(1, 16, stride=1, padding=1),
                                      nn.MaxPool2d(2,2),
                                      Flatten(),
                                      nn.Linear(1*14*14, 10))

  def forward(self, x):
    return self.net(x)